import torch
from model.vae import VAE
from model.utils import *
from model.loss_func import loss as loss
from data.load_data import CustomDataset
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler

PICTURE_SIZE=128 
batch_size=1
epoch=4000

transform = transforms.Compose([
    transforms.Resize((PICTURE_SIZE, PICTURE_SIZE)),  # 调整图片大小
    transforms.ToTensor(),           # 将图片转换为Tensor
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 归一化
])

root_dir = r'data/chicken'  # 图片文件夹路径
dataset = CustomDataset(root_dir, transform=transform)
indices = list(range(10))
sampler = SubsetRandomSampler(indices)
# 创建 DataLoader，并传入自定义的 Sampler
dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler)

device='cuda' if torch.cuda.is_available() else 'cpu'
model = VAE().to(device)
optimizer=torch.optim.AdamW(model.parameters(),lr=0.0001)

# TODO:应该划分数据集
train(net=model, optimizer=optimizer, train_loader=dataloader, loss_fc=loss, epoch=epoch)
#test(net=model, test_loader=dataloader,loss_fc=loss)
#generate(pth=)